import sqlite3
import streamlit as st
from datetime import datetime
import bcrypt  # 用于密码哈希
import logging
import os

# 获取当前路径
curr_path = os.getcwd()
db_path = os.path.join(curr_path, 'database')
log_path = os.path.join(curr_path, 'log')
# 设置日志记录
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 检查目录是否存在，如果不存在则创建它
if not os.path.exists(db_path):
    os.makedirs(db_path)
if not os.path.exists(log_path):
    os.makedirs(log_path)

# 创建一个文件处理器并设置日志级别
file_handler = logging.FileHandler(os.path.join(log_path, 'chat_app.log'))
file_handler.setLevel(logging.INFO)

# 创建一个格式化器并将其添加到处理器
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)

# 将处理器添加到logger
logger.addHandler(file_handler)

# 测试日志记录
logger.info("Logging configuration completed.")

chat_db_path = os.path.join(db_path, "chat.db")

# 初始化数据库
def init_db():
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()

        # 创建用户表（如果不存在）
        c.execute(
            """
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password TEXT NOT NULL,
                role TEXT NOT NULL DEFAULT 'user',
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
            """
        )

        # 创建会话表（如果不存在）
        c.execute(
            """
            CREATE TABLE IF NOT EXISTS conversations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                name TEXT NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users(id)
            )
            """
        )

        # 创建消息表（如果不存在）
        c.execute(
            """
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id INTEGER NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                image_path TEXT,  -- 确保包含 image_path 字段
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
            """
        )

        # 添加索引
        c.execute("CREATE INDEX IF NOT EXISTS idx_user_id ON conversations(user_id)")
        c.execute("CREATE INDEX IF NOT EXISTS idx_conversation_id ON messages(conversation_id)")

        # 初始化后注册 assistant 用户
        register_assistant(c)

        conn.commit()
        logger.info("Database initialized.")

def register_assistant(cursor):
    assistant_username = "assistant"
    assistant_password = "secure_password"  # 应该使用更安全的方式生成或配置此密码
    assistant_role = "assistant"

    try:
        # 对密码进行哈希处理
        hashed_password = bcrypt.hashpw(assistant_password.encode("utf-8"), bcrypt.gensalt())
        cursor.execute(
            "INSERT INTO users (username, password, role) VALUES (?, ?, ?)",
            (assistant_username, hashed_password, assistant_role),
        )
        logger.info(f"Assistant user '{assistant_username}' registered.")
    except sqlite3.IntegrityError:
        logger.info(f"Assistant user '{assistant_username}' already exists.")

# 添加用户（密码哈希存储）
def add_user(username, password, role="user"):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        try:
            # 对密码进行哈希处理
            hashed_password = bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt())
            c.execute(
                "INSERT INTO users (username, password, role) VALUES (?, ?, ?)",
                (username, hashed_password, role),
            )
            conn.commit()
            st.success("用户注册成功！")
        except sqlite3.IntegrityError:
            st.error("用户名已存在！")

# 验证用户（密码哈希验证）
def verify_user(username, password):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute("SELECT id, password, role FROM users WHERE username = ?", (username,))
        result = c.fetchone()
        if result and bcrypt.checkpw(password.encode("utf-8"), result[1]):
            return {"id": result[0], "username": username, "role": result[2]}
        return None

# 获取所有用户（管理员用）
def get_all_users():
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute("SELECT id, username, role, created_at FROM users")
        return c.fetchall()

# 添加新会话
def add_conversation(user_id, name):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute(
            "INSERT INTO conversations (user_id, name) VALUES (?, ?)",
            (user_id, name),
        )
        conn.commit()
        return c.lastrowid

# 更新会话时间
def update_conversation_timestamp(conversation_id):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute(
            "UPDATE conversations SET updated_at = CURRENT_TIMESTAMP WHERE id = ?",
            (conversation_id,),
        )
        conn.commit()

# 删除对话及其关联的消息
def delete_conversation(conversation_id):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()

        # 删除关联的消息
        c.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))

        # 删除对话
        c.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))

        conn.commit()

# 添加消息
def add_message(conversation_id, role, content, image_path=None):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute(
            "INSERT INTO messages (conversation_id, role, content, image_path) VALUES (?, ?, ?, ?)",
            (conversation_id, role, content, image_path),
        )
        conn.commit()

# 获取用户的所有会话
def get_conversations(user_id):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute(
            "SELECT id, name, created_at, updated_at FROM conversations WHERE user_id = ? ORDER BY updated_at DESC",
            (user_id,),
        )
        return c.fetchall()

# 获取某个会话的所有消息
def get_messages(conversation_id):
    with sqlite3.connect(chat_db_path, check_same_thread=False) as conn:
        c = conn.cursor()
        c.execute(
            "SELECT role, content, image_path, created_at FROM messages WHERE conversation_id = ? ORDER BY created_at ASC",
            (conversation_id,),
        )
        return c.fetchall()

# 初始化数据库（如果不存在）
init_db()